package geppetto.cat.programs;


import geppetto.cat.alignments.AlignmentEvaluator;
import geppetto.cat.alignments.AlignmentStats;
import geppetto.cat.alignments.AlignmentSymetrization;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.AlignmentEvaluator.Evaluation;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AgreementHMM;
import geppetto.cat.models.HMM;
import geppetto.cat.models.SubstochasticHMM;
import geppetto.phraseHMM.BijectiveHMM;
import geppetto.phraseHMM.RegularHMM;
import geppetto.phraseHMM.SymmetricHMM;

import java.io.IOException;
import java.util.Random;



/**
 * class to compute alignment error statistics (AER, Precision, Recall & their
 * derivatives
 */
public class ComputeAlignmentError {
	private String corpusDescription;

	private int size, maxSentenceSize;

	String modelDir;

	String modelName;

	String decodingType;

	float posteriorF, posteriorB;

	boolean projectTestTime; 
	private AbstractModel mhmm, mhmmb;

	private BilingualCorpus forwardCorpus, backwardCorpus;

	double slack, epsilon;
	int maxStep, maxProjectionIterations;
	
	public ComputeAlignmentError(String[] args) {
		corpusDescription = args[0];
		size = Integer.parseInt(args[1]);
		maxSentenceSize = Integer.parseInt(args[2]);
		modelDir = args[3];
		modelName = args[4];
		projectTestTime = Boolean.parseBoolean(args[5]);
		decodingType = args[6];
		if (args.length > 8){
			posteriorF = Float.parseFloat(args[7]);
			posteriorB = Float.parseFloat(args[8]);
		}else{
			posteriorF = Float.NaN;
			posteriorB = Float.NaN;
		}
		if(modelName.equalsIgnoreCase("bijective") || modelName.equalsIgnoreCase("symmetric")){
			if (args.length < 11){
				System.out.println("Missing projection parameteres");
				System.out.println("Param: corpusFile Size MaxSentenceSize modelDir modelName ProjectAtDecode decodeType posteriorValue Epsilon Slack MaxProjectionStep MaxProjectionIteratiosn ");
			}
			epsilon = Double.parseDouble(args[9]);
			slack = Double.parseDouble(args[10]);
			maxStep = Integer.parseInt(args[11]);
			maxProjectionIterations = Integer.parseInt(args[12]);
		}
		System.out.println("Corpus " + corpusDescription);
		System.out.println("Size " + size);
		System.out.println("ModelDir " + modelDir);
		System.out.println("Model Name " + modelName);
		System.out.println("ProjectAtTest Name " + projectTestTime);
		System.out.println("Decoding type " + decodingType);
		System.out.println("---");
	}

	public void initializeCorpus(String corpusDesc, int size, int maxSentSize)
			throws IOException {
		System.out.println("Loading Corpus");
		forwardCorpus = BilingualCorpus.getCorpusFromFileDescription(
				corpusDesc, size, maxSentSize);
		backwardCorpus = forwardCorpus.reverse();
		modelDir = modelDir + "/" + modelName + "/" + forwardCorpus.getName()
				+ "/" + size + "/model/";
		System.out.println("Finished Loading Corpus");
	}

	public void initializeModels(String modelName) {
		System.out.println("Loading Model");
		if (modelName.equalsIgnoreCase("baseline")) {
			mhmm = HMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward");
			mhmmb = HMM.loadModel(backwardCorpus, modelDir + "/MHMM/backward");
		} else if (modelName.equalsIgnoreCase("substochastic")) {
			mhmm = SubstochasticHMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward");
			mhmmb = SubstochasticHMM.loadModel(backwardCorpus, modelDir + "/MHMM/backward");
		} else if (modelName.equalsIgnoreCase("agreement")) {
			mhmm = AgreementHMM.loadModel(forwardCorpus, backwardCorpus,
					modelDir + "/MHMM/");
			mhmmb = ((AgreementHMM) mhmm).getBackwardModel();
		
		}else if (modelName.equalsIgnoreCase("regularHMM")) {
			mhmm = RegularHMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward");
			mhmmb = RegularHMM.loadModel(backwardCorpus, modelDir + "/MHMM/backward");			
		}else if (modelName.equalsIgnoreCase("symmetric")) {
			mhmm = SymmetricHMM.loadModel(forwardCorpus, backwardCorpus, modelDir + "/MHMM/",epsilon,slack,maxStep,maxProjectionIterations);
			mhmmb = ((SymmetricHMM) mhmm).getBackwardModel(); 
		}else if (modelName.equalsIgnoreCase("bijective")) {	
			mhmm = BijectiveHMM.loadModel(forwardCorpus, modelDir + "/MHMM/forward",epsilon,slack,maxStep,maxProjectionIterations);
			mhmmb = BijectiveHMM.loadModel(backwardCorpus, modelDir + "/MHMM/backward",epsilon,slack,maxStep,maxProjectionIterations);
			System.out.println("slack " + ((BijectiveHMM) mhmm)._slack);
			System.out.println("epsilon " + ((BijectiveHMM) mhmm)._epsilon);
			System.out.println("max steps " + ((BijectiveHMM) mhmm)._maxStepSize);
			System.out.println("max number iterartions " + ((BijectiveHMM) mhmm)._maxNumberIterations);
		} else {
			System.out.println("Unknown Model");
			System.exit(1);
		}
		System.out.println("Finished Loading Model");
	}

	public void run() throws IOException {
		initializeCorpus(corpusDescription, size, maxSentenceSize);
		initializeModels(modelName);
		AlignmentsSet align_uvf = null, align_uvb = null;
		if (decodingType.equalsIgnoreCase("viter")) {
			align_uvf = mhmm.viterbiAlignments(BilingualCorpus.TEST_CORPUS,projectTestTime);
			align_uvb = mhmmb.viterbiAlignments(BilingualCorpus.TEST_CORPUS,projectTestTime);
		} else if (decodingType.equalsIgnoreCase("post-aer")) {
			float forwardThresh = mhmm
					.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,projectTestTime);
			float backwardThresh = mhmmb
					.tuneTreshholdAER(BilingualCorpus.DEV_CORPUS,projectTestTime);
			align_uvf = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					forwardThresh,projectTestTime,true);
			align_uvb = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					backwardThresh, projectTestTime,true);
		} else if (decodingType.equalsIgnoreCase("post-f1")) {
			float forwardThresh = mhmm.tuneTreshholdF1(
					BilingualCorpus.DEV_CORPUS, posteriorF,projectTestTime);
			float backwardThresh = mhmmb.tuneTreshholdF1(
					BilingualCorpus.DEV_CORPUS, posteriorB,projectTestTime);
			align_uvf = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					forwardThresh,projectTestTime,true);
			align_uvb = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					backwardThresh,projectTestTime,true);

		} else if (decodingType.startsWith("post-pts")) {
			// sample ~200 sentences from the training corpus
			System.out.println("Forward" + mhmm);
			System.out.println("Backward" + mhmmb);
			Random r = new Random(-9119629771948792409L);
			int totNumSent = forwardCorpus.getNumberOfTrainingSentences();
			int[] sample = new int[200];
			for (int i = 0; i < sample.length; i++) {
				sample[i] = r.nextInt(totNumSent);
			}
			// count number of aligned points with viterbi
			int numAlignmentsF = mhmm.viterbiAlignments(sample,
					BilingualCorpus.TRAIN_CORPUS).numberOfAlignedPoints();
			int numAlignmentsB = mhmmb.viterbiAlignments(sample,
					BilingualCorpus.TRAIN_CORPUS).numberOfAlignedPoints();
			// tune posterior to get that.
			float forwardThresh = mhmm.tuneThresholdNumPoints(numAlignmentsF,
					sample, BilingualCorpus.TRAIN_CORPUS,projectTestTime);
			float backwardThresh = mhmmb.tuneThresholdNumPoints(numAlignmentsB,
					sample, BilingualCorpus.TRAIN_CORPUS, projectTestTime);
			align_uvf = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					forwardThresh, projectTestTime,true);
			align_uvb = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					backwardThresh,projectTestTime,true);
		} else if (decodingType.equalsIgnoreCase("post-fix")) {
			align_uvf = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					posteriorF, projectTestTime,true);
			align_uvb = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					posteriorB, projectTestTime,true);
		}else if (decodingType.equalsIgnoreCase("post-fixPrecision")) {
			float forwardThresh =
				mhmm.tuneTreshholdPrecisionRecallFixOtherValue(BilingualCorpus.TEST_CORPUS, true, posteriorF, 1.E-4, projectTestTime);
			float backwardThresh =
				mhmmb.tuneTreshholdPrecisionRecallFixOtherValue(BilingualCorpus.TEST_CORPUS, true, posteriorB, 1.E-4, projectTestTime);
			align_uvf = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					forwardThresh,projectTestTime,true);
			align_uvb = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					backwardThresh,projectTestTime,true);

		}else if (decodingType.equalsIgnoreCase("post-fixRecall")) {
			float forwardThresh =
				mhmm.tuneTreshholdPrecisionRecallFixOtherValue(BilingualCorpus.TEST_CORPUS, false, posteriorF, 0.05, projectTestTime);
			float backwardThresh =
				mhmmb.tuneTreshholdPrecisionRecallFixOtherValue(BilingualCorpus.TEST_CORPUS, false, posteriorB, 0.01, projectTestTime);
			align_uvf = mhmm.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					forwardThresh,projectTestTime,true);
			align_uvb = mhmmb.posteriorAlignments(BilingualCorpus.TEST_CORPUS,
					backwardThresh,projectTestTime,true);

		}else {
			System.out.println("Unknowon decoding");
			System.exit(1);
		}
		Evaluation evalForward = AlignmentEvaluator.evaluate(align_uvf,
				forwardCorpus.getGold());
		System.out.print("Forward ");
		AlignmentStats.printPhrases(align_uvf, false, false);
		System.out.println(evalForward.toString());
		System.out.println(evalForward.toSimpleString());
		Evaluation[] evalForwardBreak = AlignmentEvaluator.evaluateRareWords(align_uvf, forwardCorpus.getGold(), forwardCorpus);
		System.out.println("Forward Rare :" +evalForwardBreak[0]);
		System.out.println(evalForwardBreak[0].toSimpleString());
		System.out.println("Forward Common :" +evalForwardBreak[1]);
		System.out.println(evalForwardBreak[1].toSimpleString());
		System.out.println("\n\n");
		
		Evaluation evalBackward = AlignmentEvaluator.evaluate(align_uvb,
				backwardCorpus.getGold());
		System.out.print("Backward ");
		AlignmentStats.printPhrases(align_uvb, false, false);
		System.out.println(evalBackward.toString());
		System.out.println(evalBackward.toSimpleString());
		Evaluation[] evalBackwardBreak = AlignmentEvaluator.evaluateRareWords(align_uvb, backwardCorpus.getGold(), backwardCorpus);
		System.out.println("Backward Rare :" +evalBackwardBreak[0]);
		System.out.println(evalBackwardBreak[0].toSimpleString());
		System.out.println("Backward Common :" +evalBackwardBreak[1]);
		System.out.println(evalBackwardBreak[1].toSimpleString());
		System.out.println("\n\n");
		
		AlignmentsSet intersection = AlignmentSymetrization.intersect(align_uvf, align_uvb);
		Evaluation evalInter = AlignmentEvaluator.evaluate(intersection, forwardCorpus
				.getGold());
		System.out.print("Intersection ");
		AlignmentStats.printPhrases(intersection, false, false);
		System.out.println(evalInter.toString());
		System.out.println(evalInter.toSimpleString());
		Evaluation[] evalIntersectionBreak = AlignmentEvaluator.evaluateRareWords(intersection, forwardCorpus.getGold(), forwardCorpus);
		System.out.println("Intersection Rare :" +evalIntersectionBreak[0]);
		System.out.println(evalIntersectionBreak[0].toSimpleString());
		System.out.println("Intersection Common :" +evalIntersectionBreak[1]);
		System.out.println(evalIntersectionBreak[1].toSimpleString());
		System.out.println("\n\n");
		
		AlignmentsSet gdf = AlignmentSymetrization.growDiagFinal(align_uvf,
				align_uvb);
		Evaluation evalGDF = AlignmentEvaluator.evaluate(gdf, forwardCorpus
				.getGold());
		System.out.print("GDF ");
		AlignmentStats.printPhrases(gdf, false, false);
		System.out.println(evalGDF.toString());
		System.out.println(evalGDF.toSimpleString());
		
		Evaluation[] evalGDFBreak = AlignmentEvaluator.evaluateRareWords(gdf, forwardCorpus.getGold(), forwardCorpus);
		System.out.println("GDF Rare :" +evalGDFBreak[0]);
		System.out.println(evalGDFBreak[0].toSimpleString());
		System.out.println("GDF Common :" +evalGDFBreak[1]);
		System.out.println(evalGDFBreak[1].toSimpleString());
		
		System.out.println("\n\n");
	}

	public static void main(String[] args) throws IOException {
		ComputeAlignmentError exp = new ComputeAlignmentError(args);
		exp.run();
	}

}
